import glob
import os
from PyPDF2 import PdfReader, PdfWriter
import concurrent.futures
from pdf2image import convert_from_path
from PIL import Image
from pytesseract import image_to_string
import re
import datetime

output_dir = os.path.join(os.getcwd(), "output")
input_dir = os.path.join(os.getcwd(), "input")
log_dir = os.path.join(os.getcwd(), "log")

try:
    os.makedirs(output_dir, exist_ok=True)
except FileExistsError:
    print("output_dir already exists")

def log(text):
    try:
        os.makedirs(log_dir, exist_ok=True)
    except FileExistsError:
        print("log_dir already exists")
    log_file = open(os.path.join(log_dir, f"{datetime.date.today()}.txt"), "a", encoding="utf-8")
    log_file.write(f"{datetime.datetime.now()}\t{text}\n")


def extract_text(pdf_reader: PdfReader, pdf_name: str, images: list[Image.Image], page_num: int):
    start_time = datetime.datetime.now().timestamp()
    output_sub_dir = os.path.join(output_dir, pdf_name)
    try:
        os.makedirs(output_sub_dir, exist_ok=True)
    except FileExistsError:
        print(f"{output_sub_dir} already exists")
    img = images[page_num].convert('L')
    # img.save(os.path.join(output_sub_dir, f"{page_num}.png"), "PNG")
    custom_config = '--oem 3 --psm 6'
    text = image_to_string(img, lang="chi_sim", config=custom_config)
    img.close()
    text = text.replace(" ", "")
    # open(os.path.join(output_sub_dir, f"{page_num}.txt"), "w", encoding="utf-8").write(text)
    pattern = re.compile(r"检验编号[:|：|，|,]([A-Z]+[0-9]+)")
    matches = pattern.findall(text)
    page_name = f"{page_num}"
    success = False
    if len(matches) > 0:
        page_name = matches[0]
        success = True
    pdf_writer = PdfWriter()
    pdf_obj = pdf_reader.pages[page_num]
    pdf_writer.add_page(pdf_obj)
    with open(f"{output_sub_dir}/{page_name}.pdf", "wb") as output_file:
        pdf_writer.write(output_file)
    log(f"{pdf_name}-{page_num}\t是否成功：{success}\t耗时：{datetime.datetime.now().timestamp() - start_time}\t")

def splitPdfFile(pdf_path):
    with open(pdf_path, "rb") as pdf_file:
        pdf_name = pdf_file.name.split(os.path.sep).pop().split(".")[0]
        images = convert_from_path(pdf_path)
        pdf_reader = PdfReader(pdf_file)
        page_count = len(pdf_reader.pages)
        # 遍历PDF文件的每一页，将每一页保存为一个新的PDF文件
        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
            features = [executor.submit(extract_text, pdf_reader, pdf_name, images, page_num) for page_num in range(page_count)]
        for feature in features:
            feature.result()


if __name__ == "__main__":
    # 获取所有PDF文件
    curdir = os.getcwd()
    pdf_files = glob.glob(os.path.join(input_dir, "*.pdf"))
    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
      features = [executor.submit(splitPdfFile, pdfFile) for pdfFile in pdf_files]
      for feature in features:
          feature.result()
